from sklearn.datasets import load_wine
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier, export_graphviz
import numpy as np
import os

np.random.seed(0)  # 为了保证每次随机实验的结果一致

# 获取训练集、测试集
wines = load_wine()
X, Y = wines.data, wines.target
X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.2)

def decisionTreeExperiment(criterion = "gini"):
    # 获取决策树模型
    model = DecisionTreeClassifier(criterion=criterion)
    model.fit(X_train, y_train)

    # 可视化生成的决策树
    dot_file = f"DecisionTreeWithCriterion={criterion}.dot"
    export_graphviz(model, out_file=dot_file)

    p = os.path.join(os.getcwd(),dot_file)
    commandText = f"dot -Tpng {p} -o {dot_file[:-4]}.png"  # 注意，需要去 Graphviz 的官方网站下载适用于你操作系统的安装程序
    print('命令终端执行如下命令可获取可视化的决策树效果：\n',commandText)

    # 模型评估
    trainScore = model.score(X_train, y_train)
    testScore = model.score(X_test, y_test)
    print("DecisionTree experiment with 'criterion={}' ".format(criterion))
    print(f'\t trainScore:{trainScore} \n\t testScore:{testScore}\n')

dict_algorithm2criterion = {'ID3':'entropy', 'C4.5':'log_loss', "CART":'gini'}
for k, criterion in dict_algorithm2criterion.items():
    print(k)
    decisionTreeExperiment(criterion)